package com.gitee.xmhzzz.mars.code.mqtt.broker.handle;

import com.gitee.xmhzzz.mars.code.mqtt.broker.domain.MqttClient;
import io.netty.channel.ChannelHandlerContext;
import io.netty.channel.SimpleChannelInboundHandler;
import io.netty.handler.codec.mqtt.*;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.Date;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;

import static io.netty.handler.codec.mqtt.MqttMessageType.PINGRESP;
import static io.netty.handler.codec.mqtt.MqttQoS.AT_MOST_ONCE;

/**
 * @ClassName MqttBrokerHandler
 * @Description
 * @Author wzq
 * @Date 2024/9/29 15:09
 * @Version 1.0
 */
public class MqttBrokerHandler extends SimpleChannelInboundHandler<MqttMessage> {

    private static final Logger logger = LoggerFactory.getLogger(MqttBrokerHandler.class);

    private static final Map<String, String> userPasswords = new HashMap<>();
    private static final Map<String, Integer> failedAttempts = new HashMap<>();
    private static final long LOCKOUT_DURATION = 5 * 60 * 1000; // 5 minutes in milliseconds

    private static final Map<String, MqttClient> clients = new HashMap<>();

    static {
        // 假设这里有一些预设的用户名和密码
        userPasswords.put("user1", "password1");
        userPasswords.put("user2", "password2");
    }

    @Override
    protected void channelRead0(ChannelHandlerContext ctx, MqttMessage msg) {
        if (msg instanceof MqttConnectMessage) {
            handleConnect(ctx, (MqttConnectMessage) msg);
        } else if (msg instanceof MqttPublishMessage) {
            handlePublish(ctx, (MqttPublishMessage) msg);
        } else if (msg instanceof MqttSubscribeMessage) {
            handleSubscribe(ctx, (MqttSubscribeMessage) msg);
        } else if (msg instanceof MqttUnsubscribeMessage) {
            handleUnsubscribe(ctx, (MqttUnsubscribeMessage) msg);
        } else if (msg.fixedHeader().messageType().equals(MqttMessageType.PINGREQ) ) {
            handlePingReq(ctx, msg);
        } else if (msg.fixedHeader().messageType().equals(MqttMessageType.DISCONNECT)) {
            handleDisconnect(ctx,  msg);
        } else {
            logger.warn("Unknown MQTT message type: {}", msg.getClass().getName());
        }
    }

    private void handleConnect(ChannelHandlerContext ctx, MqttConnectMessage msg) {
        String username = msg.payload().userName();
        String password = msg.payload().passwordInBytes()!= null? new String(msg.payload().passwordInBytes()) : null;

        if (userPasswords.containsKey(username) && userPasswords.get(username).equals(password)) {
            // 验证通过，记录客户端信息
            MqttClient client = new MqttClient(username, ctx.channel().id().toString(), null, ctx.channel().remoteAddress().toString());
            clients.put(client.getClientId(), client);

            MqttConnAckMessage connAck = MqttMessageBuilders.connAck()
                    .sessionPresent(false)
                    .returnCode(MqttConnectReturnCode.CONNECTION_ACCEPTED)
                    .build();
            ctx.writeAndFlush(connAck);
        } else {
            // 验证失败，记录失败次数
            failedAttempts.put(username, failedAttempts.getOrDefault(username, 0) + 1);
            if (failedAttempts.get(username) >= 3) {
                // 超过三次失败，锁定账号
                failedAttempts.put(username, 0);
                // 这里可以实现锁定账号的逻辑，例如记录锁定时间
                // lockedAccounts.put(username, new Date().getTime() + LOCKOUT_DURATION);
            }

            MqttConnAckMessage connAck = MqttMessageBuilders.connAck()
                    .sessionPresent(false)
                    .returnCode(MqttConnectReturnCode.CONNECTION_REFUSED_BAD_USER_NAME_OR_PASSWORD)
                    .build();
            ctx.writeAndFlush(connAck);
            ctx.close();
        }
    }

    private void handlePublish(ChannelHandlerContext ctx, MqttPublishMessage msg) {
        // 处理发布消息
        // 这里可以根据需要实现消息的路由和分发逻辑
        logger.info("Received PUBLISH message: {}", msg);

        // 获取发布消息的主题
        String topic = msg.variableHeader().topicName();

        // 假设这里有一个方法可以获取所有已连接的客户端
        List<MqttClient> allClients = getAllClients();

        // 遍历所有客户端，检查是否有匹配的订阅
        for (MqttClient client : allClients) {
            List<String> subscriptions = client.getSubscriptions();
            if (subscriptions.stream().anyMatch(subscription -> subscription.equals(topic) ||
                    subscriptionMatches(subscription, topic))) {
                // 找到了匹配的订阅，将消息发送给该客户端
                ctx.writeAndFlush(msg);
                break; // 停止遍历当前客户端的订阅，因为已经找到了匹配
            }
        }
    }

    private boolean subscriptionMatches(String subscription, String topic) {
    // 单层通配符匹配
    if (subscription.contains("+")) {
        String regex = subscription.replace("+", "[^/]+");
        return topic.matches(regex);
    }
    // 多层通配符匹配
    else if (subscription.contains("#")) {
        String regex = subscription.replace("#", ".*");
        return topic.matches(regex);
    }
    // 不包含通配符，直接比较
    else {
        return subscription.equals(topic);
    }
}



    private List<MqttClient> getAllClients(){
        // 假设这里有一个方法可以获取所有已连接的客户端
        return clients.values().stream().collect(Collectors.toList());

    }

    private void handleSubscribe(ChannelHandlerContext ctx, MqttSubscribeMessage msg) {
        // 处理订阅请求
        MqttSubAckMessage subAck = MqttMessageBuilders.subAck()
                .addGrantedQoses(MqttQoS.AT_LEAST_ONCE, MqttQoS.AT_LEAST_ONCE)
                .build();
        ctx.writeAndFlush(subAck);

        // 记录客户端订阅的主题
        String clientId = ctx.channel().id().asShortText();
        MqttClient client = clients.get(clientId);
        if (client!= null) {
            List<String> topics = msg.payload()
                    .topicSubscriptions()
                    .stream()
                    .map(MqttTopicSubscription::topicName)
                    .collect(Collectors.toList());
            client.setSubscriptions(topics);
        }
    }

    private void handleUnsubscribe(ChannelHandlerContext ctx, MqttUnsubscribeMessage msg) {
        // 处理取消订阅请求
        MqttUnsubAckMessage unsubAck = MqttMessageBuilders.unsubAck()
                .build();
        ctx.writeAndFlush(unsubAck);

        // 移除客户端订阅的主题
        String clientId = ctx.channel().id().asShortText();
        MqttClient client = clients.get(clientId);
        if (client!= null) {
            List<String> topics = msg.payload().topics();
            client.getSubscriptions().removeAll(topics);
        }
    }

    private void handlePingReq(ChannelHandlerContext ctx, MqttMessage msg) {
        // 处理心跳请求
        MqttMessage pingResp = new MqttMessage(new MqttFixedHeader(PINGRESP, false, AT_MOST_ONCE, false, 0));
        ctx.writeAndFlush(pingResp);
    }

    private void handleDisconnect(ChannelHandlerContext ctx, MqttMessage msg) {
        // 处理断开连接请求
        String clientId = ctx.channel().id().asShortText();
        MqttClient client = clients.remove(clientId);
        if (client!= null) {
            // 可以在这里记录客户端的断开连接时间等信息
            client.setDisconnectTime(new Date());
        }
        ctx.close();
    }

    @Override
    public void exceptionCaught(ChannelHandlerContext ctx, Throwable cause) throws Exception {
        logger.error("Exception caught in MqttBrokerHandler", cause);
        ctx.close();
    }
}


